//
// Created by iris on 24-11-3.
//

#define DRIVER_MAIN

#include <fstream>
#include <iostream>
#include <sstream>
#include "AST/AST.h"
#include "ASTVisitor/MLIRASTVisitor.h"
#include "CodeGen/CodeGenerator.h"
#include "llvm/Pass.h"
#include "llvm/CodeGen/CommandFlags.h"
#include "llvm/IR/IRPrintingPasses.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/MC/TargetRegistry.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/InitLLVM.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/raw_os_ostream.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/ThreadPool.h"
#include "llvm/Support/ToolOutputFile.h"
#include "llvm/Support/WithColor.h"
#include "llvm/TargetParser/Host.h"
#include "llvm/TargetParser/Triple.h"
#include "misc/ANTLR_interface.h"

using namespace std;
using namespace llvm;

static llvm::codegen::RegisterCodeGenFlags CGF;

static llvm::cl::opt<std::string>
InputFile(llvm::cl::Positional,
          llvm::cl::desc("<input-files>"),
          cl::init("-"));

static llvm::cl::opt<std::string>
OutputFilename("o",
               llvm::cl::desc("Output filename"),
               llvm::cl::value_desc("filename"));

static llvm::cl::opt<std::string> MTriple(
    "mtriple",
    llvm::cl::desc("Override target triple for module"));

static llvm::cl::opt<bool> EmitLLVM(
    "emit-llvm",
    llvm::cl::desc("Emit IR code instead of assembler"),
    llvm::cl::init(false));

static llvm::cl::opt<bool> EmitMLIR(
    "emit-mlir",
    llvm::cl::desc("Emit MLIR code for debug usage"),
    llvm::cl::init(false)
);

static llvm::cl::opt<bool> EmitLLVMDialect(
    "emit-llvm-dialect",
    llvm::cl::desc("Emit LLVM dialect code for debug usage"),
    llvm::cl::init(false)
);


llvm::TargetMachine *
createTargetMachine(const char *Argv0) {
    llvm::Triple Triple = llvm::Triple(
        !MTriple.empty()
            ? llvm::Triple::normalize(MTriple)
            : llvm::sys::getDefaultTargetTriple());

    llvm::TargetOptions TargetOptions =
            llvm::codegen::InitTargetOptionsFromCodeGenFlags(Triple);
    std::string CPUStr = llvm::codegen::getCPUStr();
    std::string FeatureStr = llvm::codegen::getFeaturesStr();

    std::string Error;
    const llvm::Target *Target =
            llvm::TargetRegistry::lookupTarget(
                llvm::codegen::getMArch(), Triple, Error);

    if (!Target) {
        llvm::WithColor::error(llvm::errs(), Argv0) << Error;
        return nullptr;
    }

    llvm::TargetMachine *TM = Target->createTargetMachine(
        Triple.getTriple(), CPUStr, FeatureStr, TargetOptions,
        std::optional<llvm::Reloc::Model>(
            llvm::codegen::getRelocModel()));
    return TM;
}

bool emit(StringRef Argv0, llvm::Module *M,
          llvm::TargetMachine *TM,
          StringRef InputFilename) {
    CodeGenFileType FileType = llvm::codegen::getFileType();
    if (OutputFilename.empty()) {
        if (InputFilename == "-") {
            OutputFilename = "-";
        } else {
            if (InputFilename.ends_with(".mod"))
                OutputFilename =
                        InputFilename.drop_back(4).str();
            else
                OutputFilename = InputFilename.str();
            switch (FileType) {
                case llvm::CodeGenFileType::AssemblyFile:
                    OutputFilename.append(EmitLLVM ? ".ll" : ".s");
                    break;
                case llvm::CodeGenFileType::ObjectFile:
                    OutputFilename.append(".o");
                    break;
                case llvm::CodeGenFileType::Null:
                    OutputFilename.append(".null");
                    break;
            }
        }
    }

    // Open the file.
    std::error_code EC;
    sys::fs::OpenFlags OpenFlags = sys::fs::OF_None;
    if (FileType == llvm::CodeGenFileType::AssemblyFile)
        OpenFlags |= sys::fs::OF_TextWithCRLF;
    auto Out = std::make_unique<llvm::ToolOutputFile>(
        OutputFilename, EC, OpenFlags);
    if (EC) {
        WithColor::error(llvm::errs(), Argv0)
                << EC.message() << '\n';
        return false;
    }

    legacy::PassManager PM;
    if (FileType == llvm::CodeGenFileType::AssemblyFile && EmitLLVM) {
        PM.add(createPrintModulePass(Out->os()));
    } else {
        if (TM->addPassesToEmitFile(PM, Out->os(), nullptr,
                                    FileType)) {
            WithColor::error(llvm::errs(), Argv0)
                    << "No support for file type\n";
            return false;
        }
    }
    PM.run(*M);
    Out->keep();
    return true;
}

static const char *Head = "iriscc by azufre";

int main(int Argc, char **Argv) {
    // std::cout << "Hello, World!" << std::endl;
    //
    // stringstream ss;
    // ifstream fin("c_code/main.c");
    //
    // ss << fin.rdbuf();
    //
    // cout << ss.str() << endl;
    //
    // antlr4::ANTLRInputStream input(ss.str());
    // CLexer lexer(&input);
    // antlr4::CommonTokenStream tokens(&lexer);
    // CParser parser(&tokens);
    //
    // MyCVisitor cv;
    //
    // cv.visitCompilationUnit(parser.compilationUnit());

    llvm::InitLLVM X(Argc, Argv);

    InitializeAllTargets();
    InitializeAllTargetMCs();
    InitializeAllAsmPrinters();
    InitializeAllAsmParsers();

    llvm::cl::ParseCommandLineOptions(Argc, Argv, Head);

    if (llvm::codegen::getMCPU() == "help" ||
        std::any_of(llvm::codegen::getMAttrs().begin(),
                    llvm::codegen::getMAttrs().end(),
                    [](const std::string &a) {
                        return a == "help";
                    })) {
        auto Triple = llvm::Triple(LLVM_DEFAULT_TARGET_TRIPLE);
        std::string ErrMsg;
        if (auto Target = llvm::TargetRegistry::lookupTarget(
            Triple.getTriple(), ErrMsg)) {
            llvm::errs() << "Targeting " << Target->getName()
                    << ". ";
            // This prints the available CPUs and features of the
            // target to stderr...
            Target->createMCSubtargetInfo(
                Triple.getTriple(), llvm::codegen::getCPUStr(),
                llvm::codegen::getFeaturesStr());
        } else {
            llvm::errs() << ErrMsg << "\n";
            exit(EXIT_FAILURE);
        }
        exit(EXIT_SUCCESS);
    }

    llvm::TargetMachine *TM = createTargetMachine(Argv[0]);
    if (!TM)
        exit(EXIT_FAILURE);

    // todo: 不使用llvm读文件，使用原生操作
    llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer> >
            FileOrErr = llvm::MemoryBuffer::getFile(InputFile);
    if (std::error_code BufferError = FileOrErr.getError()) {
        llvm::WithColor::error(llvm::errs(), Argv[0])
                << "Error reading " << InputFile << ": "
                << BufferError.message() << "\n";
    }

    llvm::SourceMgr SrcMgr;

    // Tell SrcMgr about this buffer, which is what the
    // parser will pick up.
    SrcMgr.AddNewSourceBuffer(std::move(*FileOrErr),
                              llvm::SMLoc());

    auto s = SrcMgr.getMemoryBuffer(SrcMgr.getMainFileID())->getBuffer().str();
    // run_antlr(s);

    // todo: 插入预处理器
    // Implemented imp;
    // imp.parse(s);
    parse_func(s);

    // assert(ast::root);

    // ast_visitor::BaseASTVisitor tav(InputFile, TM);
    // tav.visitDecls(ast::root);

    ast_visitor::MLIR_ASTVisitor tav(InputFile, TM, EmitMLIR, EmitLLVMDialect);
    tav.visit(ast::root);

    tav.M->print(llvm::outs(), nullptr);

    //
    //
    // llvm::LLVMContext ctx;
    // auto cg = ::codegen::CodeGenerator::create(ctx, TM);
    //
    // auto M = cg->run(ast::root, InputFile);
    // emit(Argv[0], M.get(), TM, InputFile);
    // emit(Argv[0], tav.M, TM, InputFile);
    // emit(Argv[0], tav.M, TM, InputFile);

    // delete cg;

    return 0;
}
